#!/usr/bin/env python3

import hybrid_astar as ha
import MPC_XY_Frame as mxf
import Parameter as P
import math
import rospy
from geometry_msgs.msg import Twist
from nav_msgs.msg import Odometry
import time
import sys
import matplotlib.pyplot as plt
import copy


# 单车：仿真点到点
class OneCarNavSim:
    def __init__(self, StartPoint, globalPoint, isShow=False):
        if not isShow:
            self.sub = rospy.Subscriber("odom", Odometry, self.controlCallback, queue_size=1)
            self.pub = rospy.Publisher("cmd_vel", Twist, queue_size=1)
        self.iter = 0
        self.actX = []
        self.actY = []
        self.actYaw = []
        self.curPoint = StartPoint
        self.globalPoint = globalPoint

        # aStar
        self.aCurPoint = copy.deepcopy(StartPoint)
        self.aGlobalPoint = copy.deepcopy(globalPoint)
        self.rePlanIter = P.C.AStar.RE_PlAN_ITER
        self.aStarPlot = True   # 是否显示astar轨迹
        self.path = None
        self.astarX = []
        self.astarY = []

        # mpc
        self.dt = P.C.MPC.dt
        self.node = mxf.Node(x=self.curPoint[0], y=self.curPoint[1], yaw=self.curPoint[2], 
            v=0.0,time=round(time.time(), 2))
        self.aOpt, self.deltaOpt, self.xOpt, self.yOpt, self.yawOpt, self.vOpt = 0, 0, 0, 0, 0, 0
        self.target_ind = 0
        self.sp = None
        self.ref_path = None
        self.cx, self.cy, self.cyaw, self.ck = None, None, None, None
        self.isMPCPlot = False      # 是否打印信息
        self.isMPCPrint = True      # 是否显示小车轨迹
        self.twist = Twist()

    def controlCallback(self, data):
        # 更新车辆信息
        startTime = time.time()
        self.updateNodeFromSim(data)
        self.actX.append(self.node.x)
        self.actY.append(self.node.y)
        self.actYaw.append(self.node.yaw)

        # 轨迹重规划
        if (self.iter % self.rePlanIter) == 0:
            self.curPoint = [self.node.x, self.node.y, self.node.yaw]
            self.aStarTrajectoryPlanning()
        self.iter += 1

        # mpc
        v, delta = self.mpcControl()
        self.twist.linear.x = v  # 线速度 m/s
        self.twist.angular.z = delta  # 倾角 rad
        self.pub.publish(self.twist)

        # 终点判断
        dist = math.hypot(self.node.x - self.globalPoint[0], self.node.y - self.globalPoint[1])  # 计算终点距离
        if dist < P.C.MPC.dist_stop and self.node.v <= P.C.MPC.speed_stop:               # 终点判断
            plt.show()
            rospy.signal_shutdown("over")

        # 显示
        if self.isMPCPrint:
            print("终点距离:", dist)
            print("速度:", v, "倾角:", delta)
            print("小车信息,x:", self.node.x, "y:", self.node.y, "v:",self.node.v, "yaw:", math.degrees(self.node.yaw))

        if self.isMPCPlot:
            dy = (self.node.yaw - self.actYaw[-2]) / (self.node.v * self.dt)
            steer = mxf.rs.pi_2_pi(-math.atan(P.C.Robot.WB * dy))  # 计算前轮倾角
            plt.cla()
            mxf.draw.draw_car(self.node.x, self.node.y, self.node.yaw, steer, mxf.P)
            plt.gcf().canvas.mpl_connect('key_release_event',
                                        lambda event:[exit(0) if event.key == 'escape' else None])
            if self.xOpt is not None:
                plt.plot(self.xOpt, self.yOpt, color='darkviolet', marker='*')
                plt.plot(self.cx, self.cy, color='gray')
                plt.plot(self.actX, self.actY, '-b')
                plt.plot(self.cx[self.target_ind], self.cy[self.target_ind])
                plt.axis("equal")
                plt.title("Linear MPC, " + "v = " +
                        str(round(self.node.v * 3.6, 2)))
                plt.pause(0.001)

        # 控制发布信息评率
        endTime = time.time()
        timeSleep = self.dt - (endTime - startTime)
        # print("timeSleep:",timeSleep)
        if timeSleep > 0:
            time.sleep(timeSleep)

    # 更新车辆位置信息
    def updateNodeFromSim(self, data):
        x_cur = data.pose.pose.position.x
        y_cur = data.pose.pose.position.y
        time_cur = round(time.time(), 2)
        self.node.yaw = self.eulerAnglesToEulerAngles(
            data.pose.pose.orientation)
        self.node.v = self.getV(self.node.x, self.node.y,
                                self.node.time, x_cur, y_cur, time_cur)
        self.node.x = x_cur
        self.node.y = y_cur
        self.node.time = time_cur

    # 四元数计算车辆姿态角
    def eulerAnglesToEulerAngles(self, q):
        siny_cosp = 2 * (q.w * q.z + q.x * q.y)
        cosy_cosp = 1 - 2 * (q.y * q.y + q.z * q.z)
        yaw = math.atan2(siny_cosp, cosy_cosp)
        return yaw

    # 计数速度
    def getV(self, x1, y1, t1, x2, y2, t2):
        v = math.sqrt((x1-x2)**2 + (y1-y2)**2)/(t2-t1)
        # 最大速
        if v >= P.C.Robot.speed_max:
            v == P.C.Robot.speed_max
        # 正负判断
        yaw = math.atan2(x2-x1, y2-y1)
        if abs(yaw - self.node.yaw) >= 1.57:
            v = -v
        return v

    # astar轨迹规划
    def aStarTrajectoryPlanning(self):
        self.aStarPretreatment()
        self.path = ha.main(self.aCurPoint, self.aGlobalPoint,self.aStarPlot)   # a星规划
        self.aStarPostProcessing()

        self.cx, self.cy, self.cyaw, self.ck, s = mxf.cs.calc_spline_course(
            self.astarX, self.astarY, ds=P.C.MPC.d_dist)    # 生成参考轨迹
        self.sp = mxf.calc_speed_profile(self.cx, self.cy, self.cyaw, P.C.MPC.target_speed)   # 初略速度规划
        self.ref_path = mxf.PATH(self.cx, self.cy, self.cyaw, self.ck)                   # 参考轨迹
        self.deltaOpt, self.aOpt = None, None    # 重置控制量
        self.actX, self.actY, self.actYaw = [self.node.x], [self.node.y], [self.node.yaw]  # 重置真实轨迹

    # astar数据预处理
    def aStarPretreatment(self):
        self.aCurPoint[0] = self.curPoint[0] / ha.C.N + ha.C.MAP_SIZE[0] / 2
        self.aCurPoint[1] = self.curPoint[1] / ha.C.N + ha.C.MAP_SIZE[1] / 2
        self.aGlobalPoint[0] = self.globalPoint[0] / ha.C.N + ha.C.MAP_SIZE[0] / 2
        self.aGlobalPoint[1] = self.globalPoint[1] / ha.C.N + ha.C.MAP_SIZE[1] / 2

    # astar数据后处理
    def aStarPostProcessing(self):
        self.astarX.clear()
        self.astarY.clear()
        for i in range(len(self.path.x)):
            self.astarX.append((self.path.x[i] - ha.C.MAP_SIZE[0]/2) * ha.C.N)
            self.astarY.append((self.path.y[i] - ha.C.MAP_SIZE[1]/2) * ha.C.N)

    # MPC规划
    def mpcControl(self):
        # mpc
        z_ref, self.target_ind = mxf.calc_ref_trajectory_in_T_step(
            self.node, self.ref_path, self.sp)   # 计算预测结果
        z0 = [self.node.x, self.node.y, self.node.v, self.node.yaw]
        self.aOpt, self.deltaOpt, self.xOpt, self.yOpt, self.yawOpt, self.vOpt = mxf.linear_mpc_control(
            z_ref, z0, self.aOpt, self.deltaOpt)    # MPC

        if self.deltaOpt is not None:
            deltaExc, aExc = self.deltaOpt[0], self.aOpt[0]
            vExc = self.node.v + self.dt*aExc

        return vExc, deltaExc

    def justShow(self):
        # 轨迹规划
        self.aStarTrajectoryPlanning()

        while True:
            # mpc
            v, delta = self.mpcControl()

            # 更新节点
            self.node.update(self.aOpt[0],delta,1)
            self.actX.append(self.node.x)
            self.actY.append(self.node.y)
            self.actYaw.append(self.node.yaw)

            # 终点判断
            dist = math.hypot(
                self.node.x - self.globalPoint[0], self.node.y - self.globalPoint[1])  # 计算终点距离
            if dist < P.C.MPC.dist_stop and self.node.v <= P.C.MPC.speed_stop:               # 终点判断
                break

            ## 画图
            if self.isMPCPrint:
                print("终点距离:", dist)
                print("速度:", v, "倾角:", delta)
                print("小车信息,x:", self.node.x, "y:", self.node.y, "v:",
                    self.node.v, "yaw:", math.degrees(self.node.yaw))

            dy = (self.node.yaw - self.actYaw[-2]) / (self.node.v * self.dt)
            steer = mxf.rs.pi_2_pi(-math.atan(P.C.Robot.WB * dy))  # 计算前轮倾角
            plt.cla()
            mxf.draw.draw_car(self.node.x, self.node.y, self.node.yaw, steer, mxf.P)
            plt.gcf().canvas.mpl_connect('key_release_event',
                                        lambda event:[exit(0) if event.key == 'escape' else None])
            if self.xOpt is not None:
                plt.plot(self.xOpt, self.yOpt, color='darkviolet', marker='*')
                plt.plot(self.cx, self.cy, color='gray')
                plt.plot(self.actX, self.actY, '-b')
                plt.plot(self.cx[self.target_ind], self.cy[self.target_ind])
                plt.axis("equal")
                plt.title("Linear MPC, " + "v = " +
                        str(round(self.node.v * 3.6, 2)))
                plt.pause(0.001)

        print("over-------------------")
        plt.show()

if __name__ == '__main__':
    rospy.init_node("one_robot_astar_mpc",anonymous=False)
    StartPoint = [0, 0, math.radians(45)]
    globalPoint = [1, 1, math.radians(120)]
    isJustShow = False
    oneCarNavSim = OneCarNavSim(StartPoint, globalPoint, isJustShow)
    if isJustShow:
        oneCarNavSim.justShow()
    rospy.spin()